package org.huangrui.spark.scala.sql

import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.{Encoder, Encoders, SparkSession, functions}
import org.apache.spark.sql.functions.udaf

/**
 * @Author hr
 * @Create 2024-10-20 18:12 
 */
object SparkSQL03_SQL_UDAF {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().appName("SparkSQL03_SQL_UDAF").master("local[2]").getOrCreate()
    spark.read.json("data/user.json").createOrReplaceTempView("user")

    // TODO SparkSQL采用特殊的方式将UDAF转换成UDF使用
    //      UDAF使用时需要创建自定义聚合对象
    //        udaf方法需要传递2个参数
    //             第一个参数表示UDAF对象
    //             第二个参数表示UDAF对象对应的Encoder对象

    spark.udf.register("ageAvg", functions.udaf(new MyAvgUDAF))

    spark.sql("select ageAvg(age) from user").show()

    spark.stop()
  }

  case class Buff(var total: Long, var count: Long)

  /**
   * 自定义聚合函数类：计算年龄的平均值
   *      1. 继承org.apache.spark.sql.expressions.Aggregator, 定义泛型
   *         IN : 输入的数据类型 Long
   *         BUF : 缓冲区的数据类型 Buff
   *         OUT : 输出的数据类型 Long
   *         2. 重写方法(6)
   */
  class MyAvgUDAF extends Aggregator[Long, Buff, Long] {
    // z & zero : 初始值或零值
    // 缓冲区的初始化
    override def zero: Buff = {
      Buff(0, 0)
    }

    // 根据输入的数据更新缓冲区的数据
    override def reduce(buff: Buff, in: Long): Buff = {
      buff.total = buff.total + in
      buff.count = buff.count + 1
      buff
    }

    // 合并缓冲区
    override def merge(buff1: Buff, buff2: Buff): Buff = {
      buff1.total = buff1.total + buff2.total
      buff1.count = buff1.count + buff2.count
      buff1
    }

    //计算结果
    override def finish(buff: Buff): Long = {
      buff.total / buff.count
    }

    // 缓冲区的编码操作
    override def bufferEncoder: Encoder[Buff] = Encoders.product

    // 输出的编码操作
    override def outputEncoder: Encoder[Long] = Encoders.scalaLong
  }
}
